#!/usr/bin/env python
import bisect
import numpy as np
import matplotlib as mpl
from matplotlib import pyplot as plt
from matplotlib.patches import Polygon
from astropy import time
from astropy.io import fits as pf
from maputility2 import skyutils

t0 = time.Time('2001-01-01')

def time2met(t):
    "t should be formatted in string like 2008-08-04T00:00:00"
    return (time.Time(t)-t0).sec

def hhmmss2ddmmss(hhmmss):
    pass

def extract_sc(scfile, met0, dmet=86400, ext_name='SC_DATA', islat=True):
    """
    extract the sc state between the time interval

    Parameters
    ----------
    scfile : str
        infile of the scfile
    met0, dmet : double, int. [sec]
        the time interval you interest is [met0:met0+dmet]
    ext_name : str, int
        the extension name of the scdata in the scfile
    islat: bool, True
        whether to account for the LAT specific parameters

    Returns
    -------
    scdata : numpy.ndarray, (?, 7)
        an array of (met_start-met0, met_stop-met0, ra_zen,
                    dec_zen, ra_scz, dec_scz, scflag)
    """
    met1 = met0 + dmet # in second

    hdl = pf.open(scfile)
    hdu = hdl[ext_name]

    tstart = hdu.data.field('START')
    tstop  = hdu.data.field('STOP')
    #print(met0, tstart[0], met1, tstart[-1])
    assert met0 >= tstart[0] and met1 <= tstart[-1]
    idx0 = bisect.bisect(tstart, met0)-1
    idx1 = bisect.bisect_right(tstart, met1)+1

    ra_zenith = hdu.data.field('RA_ZENITH')[idx0:idx1]
    dec_zenith = hdu.data.field('DEC_ZENITH')[idx0:idx1]
    ra_scz = hdu.data.field('RA_SCZ')[idx0:idx1]
    dec_scz = hdu.data.field('DEC_SCZ')[idx0:idx1]

    in_saa = hdu.data.field('IN_SAA')[idx0:idx1]
    if islat:
        lat_config = hdu.data.field('LAT_CONFIG')[idx0:idx1]
        data_qual = hdu.data.field('DATA_QUAL')[idx0:idx1]
        flags = (lat_config==1)&(data_qual>0)&(~in_saa)
    else:
        flags = ~in_saa
        ra_scz += 180.
        dec_scz = -dec_scz

    scdata = np.c_[tstart[idx0:idx1]-met0, tstop[idx0:idx1]-met0,
        ra_zenith, dec_zenith, ra_scz, dec_scz, flags,
        ]
    #print(idx0, idx1)
    #print(scdata[:10, :2])
    del hdu.data
    hdl.close()
    return scdata

def srcpos(ra_src, dec_src, scdata):
    "calculate the angsep between src (in degree) and zenith, scz"
    ang_zen = skyutils.angsep((ra_src, dec_src), scdata[:, 2:4])
    ang_scz = skyutils.angsep((ra_src, dec_src), scdata[:, 4:6])
    return np.c_[scdata[:, :2], ang_zen, ang_scz, scdata[:, 6]]

def plot_scpos(posdata, outfile=None, t0str=None, simple=False):
    fg, ax = plt.subplots()

    time_range = [posdata[0, 0], posdata[-1, 1]]
    ax.plot(posdata[:, 0], posdata[:, 2], ls='-', c='C0', label='zenith', alpha=0.5)
    ax.plot(posdata[:, 0], posdata[:, 3], ls='--', c='C1', label='boresight', alpha=0.5)


    ylim = ax.get_ylim()
    #ax.plot(time_range, [111.9, 111.9], ls=':', c='C0') # earth limb, Ackermann+PRL2014
    #ax.plot(time_range, [51.8, 51.8], ls=':', c='C1', label='FoV ref') # FoV=2.4, Atwood+ApJ2009

    if not simple:
        # earth limb, Ackermann+PRL2014
        vert = [(time_range[0], ylim[1]), (time_range[0], 111.9),
               (time_range[1], 111.9), (time_range[1], ylim[1])]
        poly = Polygon(vert, facecolor='none', edgecolor='C0', alpha=0.7, hatch=r"\\", label='earth')
        ax.add_patch(poly)

        # FoV=2.4, Atwood+ApJ2009
        vert = [(time_range[0], ylim[1]), (time_range[0], 51.8),
               (time_range[1], 51.8), (time_range[1], ylim[1])]
        poly = Polygon(vert, facecolor='none', edgecolor='C1', alpha=0.7, hatch=r"//", label='FoV')
        ax.add_patch(poly)

        filt = posdata[:, 4] == 0
        if filt.any():
            has_label = False
            tstart_inval, tstop_inval = posdata[filt, 0], posdata[filt, 1]
            #print(tstart_inval, tstop_inval)
            tmin0 = tstart_inval[0]
            for i in range(tstart_inval.shape[0]):
                if i+1>=tstart_inval.shape[0]:
                    tmin00 = None
                    pass
                elif tstart_inval[i+1] == tstop_inval[i]:
                    continue
                else:
                    tmin00 = tstart_inval[i+1]
                #print(tmin0, tstop_inval[i])

                if has_label:
                    ax.axvspan(tmin0, tstop_inval[i], facecolor='0.9', edgecolor='0.9', alpha=0.7)
                else:
                    ax.axvspan(tmin0, tstop_inval[i], facecolor='0.9', edgecolor='0.9', alpha=0.7, label='bad qual')
                    has_label = True
                tmin0 = tmin00

        inval_time = posdata[1:, 0] != posdata[:-1, 1]
        if inval_time.any():
            has_label = False
            t0_inval = posdata[:-1, 1][inval_time]
            t1_inval = posdata[1:, 0][inval_time]
            for t0, t1 in zip(t0_inval, t1_inval):
                if has_label:
                    ax.axvspan(t0, t1, facecolor='none', edgecolor='0.5', alpha=0.7, hatch='.')
                else:
                    ax.axvspan(t0, t1, facecolor='none', edgecolor='0.5', alpha=0.7, hatch='.', label='no obs')
                    has_label = True

    ax.set_ylim(ylim)

    ax2 = ax.twiny()
    xlim = ax.get_xlim()
    ax2.set_xlim(xlim[0]/3600., xlim[1]/3600.)
    ax2.set_xlabel('$t-T_0$ [h]')

    if t0str is None:
        ax.set_xlabel('$t-T_0$ [s]')
    else:
        ax.set_xlabel('$t-T_0$ [s]\n($T_0$={})'.format(t0str))
    ax.set_ylabel(r'$\theta$ [deg]')
    ax.legend()
    fg.tight_layout()

    if outfile is None:
        plt.show()
    else:
        fg.savefig(outfile)

if __name__ == '__main__':
    import argparse

    parser = argparse.ArgumentParser(prefix_chars='-+',
            description='draw the position of the src w.r.t. the LAT spacecraft')
    parser.add_argument("ra", type=float, help='the RA of the src [deg]')
    parser.add_argument("dec", type=float, help='the Decl of the src [deg]')
    parser.add_argument("time", type=str, help='the starting time, should be like "2008-08-04T00:00:00"')
    parser.add_argument("-scfile", type=str, default='sc.fits', help='the filename of the scfile (default: sc.fits)')
    parser.add_argument("-dtime", type=float, default=3600., help='the time interval [s] to investigate (default: 3600.)')
    parser.add_argument("-outfile", type=str, default=None, help='the name of the output file (default: None)')
    parser.add_argument("+dampe", action='store_true', default=False, help='use the DAMPE MET (default: False)')
    parser.add_argument("+simple", action='store_true', default=False, help='simply plot the theta')
    args = parser.parse_args()

    if args.dampe:
        t0 = time.Time('2013-01-01')
    if args.scfile is None:
        mpl.use('Agg')

    print('converting time to MET ...')
    met0 = time2met(args.time)
    print('MET={}'.format(met0))
    print('extracting sc information ...')
    scdata = extract_sc(scfile=args.scfile, met0=met0, dmet=args.dtime, islat=(not args.dampe))
    print('calculating the angular separation ...')
    posdata = srcpos(args.ra, args.dec, scdata)
    print('ploting ...')
    plot_scpos(posdata, outfile=args.outfile, t0str=args.time, simple=args.simple)
